"""

Copyright 2009 Michael Seiler
Rutgers University
miseiler@gmail.com

This file is part of ConsensusCluster.

ConsensusCluster is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

ConsensusCluster is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with ConsensusCluster.  If not, see <http://www.gnu.org/licenses/>.


"""

import pycuda.gpuarray as gpuarray
import numpy as N

from km_kernels import get_km_funcs
from consensus  import upd_mat, upd_mat_atom
from transpose  import tp
from sampler    import get_bootstrap_func

import kernel_helpers as kh

BOOTSTRAP     = 1
NON_BOOTSTRAP = 1

def kmeans(ary, num_clusters = 2, repeats = 300):
    
    # Functions
    init_km, km, upd_cntrds = get_km_funcs(num_clusters, 0.001)
    bootstrap = get_bootstrap_func()
    
    # Main array
    M      = kh.aligntobsize(ary, 128)
    
    # Dims
    numsamples, numfeatures = ary.shape
    n, m       = M.shape
    bstrap_dim = n

    assert num_clusters < numsamples

    M_gpu      = gpuarray.to_gpu(M)
    centroids  = gpuarray.zeros((num_clusters,m), dtype=N.float32)
    kresults   = gpuarray.zeros((bstrap_dim,), dtype=N.int32)
    moved_flag = gpuarray.zeros((1,), dtype=N.int32)
    bsM        = gpuarray.zeros((bstrap_dim,m), dtype=N.float32)
    bsM_T      = gpuarray.zeros((m,bstrap_dim), dtype=N.float32)
    samples    = gpuarray.zeros((bstrap_dim,), dtype=N.int32)
    clustcount = gpuarray.zeros((n,n), dtype=N.int32)
    totalcount = gpuarray.zeros((n,n), dtype=N.int32)
    
    # XXX Be very careful with numsamples and bootstrap dim!
    
    if BOOTSTRAP:
        
        for _ in xrange(int(repeats)):
    
            bootstrap(bsM, M_gpu, samples, numsamples, bstrap_dim)
            tp(bsM_T, bsM)
    
            init_km(centroids, bsM_T, bstrap_dim) # Note bstrap_dim in place of numsamples
    
            while 1:
                moved_flag.fill(0)
                km(kresults, bsM_T, centroids)
                upd_cntrds(centroids, bsM_T, kresults, bstrap_dim, moved_flag) # Note bstrap_dim in place of numsamples
                if not moved_flag.get()[0]: break
    
            upd_mat_atom(clustcount, totalcount, kresults, samples) # An atomic update is necessary when bootstrapping

    if NON_BOOTSTRAP:

        tp(bsM_T, M_gpu)
        samples = gpuarray.arange((n,), dtype=N.int32)
    
        for _ in xrange(int(repeats)):
    
            init_km(centroids, bsM_T, numsamples)
    
            while 1:
                moved_flag.fill(0)
                km(kresults, bsM_T, centroids)
                upd_cntrds(centroids, bsM_T, kresults, numsamples, moved_flag)
                if not moved_flag.get()[0]: break

            upd_mat(clustcount, totalcount, kresults, samples)


    return clustcount.get()[:numsamples][:,:numsamples], totalcount.get()[:numsamples][:,:numsamples]
